# -*- coding: utf-8 -*-
from __future__ import print_function
from datetime import datetime
from ddl_platform.database.base import BaseTable
from ddl_platform.database.base import main_database
from ddl_platform.ddlib.job import JobStatus, JobTaskStatus
import pickle


class ProfiledJobTable(BaseTable):
    def __init__(self, db=main_database):
        self.db = db
        self.TABLE = 'profiledjobtb'

    def query(self):
        jobs = self.db.hvals(self.TABLE)
        decoded_jobs = [pickle.loads(j) for j in jobs]
        return decoded_jobs

    def query_waiting_jobs(self):
        jobs = self.query()
        return [j for j in jobs if j._status == JobStatus.WAITING or j._status == JobStatus.MIGRATION]

    def get(self, job_id):
        j = self.db.hget(self.TABLE, job_id)
        job_info = None
        if j:
            job_info = pickle.loads(j)
        return job_info

    def update_status(self, job_id, status, exception=None, checkpoint_path=None):
        job = self.get(job_id)
        job._status = status
        if status == JobStatus.FINISHED or status == JobStatus.EXCEPTION:
            job._completion_time = datetime.timestamp(datetime.now())
        if status == JobStatus.EXCEPTION and exception is not None:
            job._exception = exception
        if status == JobStatus.MIGRATION and checkpoint_path is not None:
            job._checkpoint_path = checkpoint_path
        encoded_job = pickle.dumps(job)
        self.db.hset(self.TABLE, job._job_id, encoded_job)

    def insert(self, job):
        encoded_job = pickle.dumps(job)
        self.db.hset(self.TABLE, job._job_id, encoded_job)
        return job._job_id

    def delete(self, job_id):
        self.db.hdel(self.TABLE, job_id)

    def clear_table(self):
        self.db.delete(self.TABLE)

def clear_table():
    sjt = ProfiledJobTable()
    sjt.clear_table()

tb =  ProfiledJobTable()
profiled_job_table = tb


def test():
    import ddl_platform.scheduler.job_info as job 
    sjt = ProfiledJobTable()
    j = job.JobInfo('/home/esetstore/repos/ddl-platform/examples/cifar10_resnet20', '/home/esetstore/repos/ddl-platform/examples/cifar10_resnet20/job_0.yaml')
    print('a new job: %s' % j._job_id)
    jid = sjt.insert(j)
    print('a job inserted: ', jid)
    jobs = sjt.query()
    print(jobs)


if __name__ == '__main__':
    clear_table()
    #test()
